Skip to content

Enable inference CUDA graphs with Shortcut MoE and wide residuals - #7772

Merged
santhnm2 merged 17 commits into
NVIDIA:mainfrom
susavlsh10:nemotron4-inference-cuda-graphs
Oct 15, 2026
Merged

Enable inference CUDA graphs with Shortcut MoE and wide residuals#7772
santhnm2 merged 17 commits into
NVIDIA:mainfrom
susavlsh10:nemotron4-inference-cuda-graphs

Conversation

@susavlsh10

@susavlsh10 susavlsh10 commented Oct 1, 2026 •
edited by santhnm2
Loading

Copy link
Copy Markdown
Contributor
  • I, the PR author, have personally reviewed every line of this PR.

What does this PR do?

Enable local, block-scoped CUDA graphs for text inference with GDP, Shortcut MoE, and wide residuals.

The change fixes GDP decode padding during graph replay and the FP32 wide-residual to BF16 RMSNorm precision boundary.

Regression tests cover configuration validation, GDP padding, RMSNorm dtype, checkpoint routing, and eager-versus-graph generation through a small Shortcut MoE model.

Dependencies

Depends on #7752 (text-only MIMO checkpoint loading), #6855 (wide residual MTP/MIMO support), and #7440 (wide residual Shortcut MoE support).

This draft currently includes a tested integration merge of the wide-residual work. We will refresh the dependency revisions and narrow the diff against main as those PRs progress.

⚠️ For major changes (either in lines of code or in its impact), please make sure to first share a design doc with the team. If you're unsure what's the best way to do so, contact @NVIDIA/mcore-oncall.

Issue tracking

For PRs from open-source community contributors:

  • New features: a linked issue is required. Please open a feature request and reference it here before submitting the PR.
  • Small updates (bug fixes, minor improvements): a linked issue is recommended and will accelerate the PR review process.

Linked issue:

Contribution process

Pre-checks

  • I have added relevant unit tests
  • I have added relevant functional tests
  • If this PR adds or changes a GPU kernel (Triton, jit_fuser/torch.compile, CUDA extension, TE or external-library dispatch, or a scatter/index accumulation), I have added or updated its bit-exact determinism test and registered it in tests/unit_tests/determinism/kernels/manifest.py (guide)
  • I have added proper typing to my code Typing guidelines
  • I have added relevant documentation
  • I have run the autoformatter.sh on my PR

Code review

Feel free to message or comment @NVIDIA/mcore-oncall to help accelerate your merge into main. The less complex your PR is, the faster it will be approved and merged!

All PRs start as draft. If you open a non-draft PR, it will be automatically converted to draft.

Step 1: Mark PR as "Ready for Review"

  1. When your PR is ready, click Ready for Review.
  2. An oncall reviewer is auto-assigned and expert reviewers are notified based on your changes.
    • Some PRs may jump straight to step 2. This is determined by .github/CODEOWNERS.

⚠️ Only mark as ready once merge-conflicts are resolved and the CI is passing.
Final Review might get declined if these requirements are not fulfilled.

Step 2: Final Review

For PRs that change megatron/core, once all expert reviewers have approved, the Final Review label is applied automatically and final reviewers are assigned.

For PRs outside megatron/core, this step is skipped.

Step 3: Approved

Once all required reviewers have approved, the Approved label is applied automatically.

Merge

Any member of mcore-engineers will be able to merge your PR.

@copy-pr-bot

copy-pr-bot Bot commented Oct 1, 2026

Copy link
Copy Markdown

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

@susavlsh10
susavlsh10 marked this pull request as ready for review October 1, 2026 15:40
@susavlsh10
susavlsh10 requested review from a team as code owners October 1, 2026 15:40
mathemakitten
mathemakitten reviewed Oct 1, 2026
View reviewed changes
Comment thread megatron/core/tensor_parallel/inference_layers.py Outdated
Comment thread megatron/core/ssm/ssm_inference.py Outdated
mathemakitten
mathemakitten reviewed Oct 1, 2026
View reviewed changes
Comment thread megatron/core/transformer/transformer_layer.py Outdated
@mathemakitten mathemakitten mentioned this pull request Oct 2, 2026
Merged
1 of 7 tasks
@susavlsh10
susavlsh10 force-pushed the nemotron4-inference-cuda-graphs branch from 92bdc80 to 753297e Compare October 3, 2026 00:39
santhnm2
santhnm2 reviewed Oct 3, 2026
View reviewed changes
Comment thread megatron/core/transformer/transformer_config.py
Comment thread megatron/training/arguments.py
Comment thread megatron/core/inference/text_generation_server/dynamic_text_gen_server/vlm_dynamic_inference.py Outdated
@lhb8125 lhb8125 mentioned this pull request Oct 5, 2026
Open
@sidsingh-nvidia

sidsingh-nvidia commented Oct 5, 2026 •
edited
Loading

Copy link
Copy Markdown
Contributor

I suggest making the block scope default for inference. It is what we have been extensively using in Nemo-RL.

@susavlsh10

susavlsh10 commented Oct 5, 2026

Copy link
Copy Markdown
Contributor Author

I suggest making the block scope default for inference. It is what we have been extensively using in Nemo-RL.

@sidsingh-nvidia should we do this in this PR?

@Connor-XY Connor-XY mentioned this pull request Oct 6, 2026
Closed
@santhnm2
santhnm2 removed request for a team October 6, 2026 16:10
@susavlsh10
susavlsh10 force-pushed the nemotron4-inference-cuda-graphs branch from 040471c to 1d5f7d4 Compare October 6, 2026 17:56
@santhnm2 santhnm2 changed the title Enable Nemotron 4 text inference CUDA graphs with Shortcut MoE and wide residuals Enable inference CUDA graphs with Shortcut MoE and wide residuals Oct 6, 2026
@santhnm2

santhnm2 commented Oct 6, 2026

Copy link
Copy Markdown
Contributor

/ok to test 1d5f7d4

@santhnm2

santhnm2 commented Oct 9, 2026

Copy link
Copy Markdown
Contributor

/ok to test 95f9b94

@santhnm2
Merge with main
c436ba3 
Signed-off-by: Keshav Santhanam <ksanthanam@nvidia.com>
@santhnm2

santhnm2 commented Oct 12, 2026

Copy link
Copy Markdown
Contributor

/ok to test c436ba3

@susavlsh10
Clean up NVLS state after Shortcut MoE CUDA graph tests
d66aacb 
Signed-off-by: Susav Shrestha <susavlsh10@gmail.com>
auto-merge was automatically disabled October 12, 2026 20:41

Head branch was pushed to by a user without write access

@santhnm2

santhnm2 commented Oct 12, 2026

Copy link
Copy Markdown
Contributor

/ok to test 27c6495

@santhnm2
santhnm2 enabled auto-merge October 12, 2026 20:49
@balasaajay balasaajay mentioned this pull request Oct 12, 2026
Open
@balasaajay

balasaajay commented Oct 14, 2026

Copy link
Copy Markdown
Contributor

/ok to test 0c70a59

@nemo-automation-bot

nemo-automation-bot Bot commented Oct 15, 2026

Copy link
Copy Markdown

🔄 Merge queue validation started!

You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/38693382074

@nemo-automation-bot

nemo-automation-bot Bot commented Oct 15, 2026

Copy link
Copy Markdown

🔄 Merge queue validation started!

You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/38705174583

@nemo-automation-bot

nemo-automation-bot Bot commented Oct 15, 2026

Copy link
Copy Markdown

🔄 Merge queue validation started!

You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/38710370155

This branch was successfully deployed

2 active deployments
public — 0c70a591 Deployed Oct 14, 2026 by copy-pr-bot[bot] via build-test-publish-wheels / publish-wheels (megatron-core, arm64) #11824
test — 0c70a591 Deployed Oct 14, 2026 by copy-pr-bot[bot] via cicd-wait-in-queue #24139
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@Mellonta Mellonta Mellonta left review comments
@mathemakitten mathemakitten mathemakitten left review comments
@kwangjunahn kwangjunahn kwangjunahn left review comments
@deepakn94 deepakn94 deepakn94 approved these changes
@santhnm2 santhnm2 santhnm2 approved these changes
@yaoyu-33 yaoyu-33 yaoyu-33 approved these changes
@sidsingh-nvidia sidsingh-nvidia sidsingh-nvidia approved these changes

Assignees

No one assigned

Labels

Approved All necessary approvals have been made

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

10 participants

@susavlsh10 @sidsingh-nvidia @santhnm2 @balasaajay @deepakn94 @Mellonta @mathemakitten @yaoyu-33 @kwangjunahn @svcnvidia-nemo-ci